import cv2
import os
import numpy as np

def separate_mask(image_id, filename):
    im = cv2.imread(filename, cv2.IMREAD_GRAYSCALE)
    # print(len(im[im==0]))
    # print(len(im[im==255]))
    retval, im = cv2.threshold(im, 100, 255, cv2.THRESH_BINARY)

    _, contours, hierarchy = cv2.findContours(im.copy(),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)

    for idx in range(len(contours)):
        mask = np.zeros_like(im)
        cv2.drawContours(mask, contours, idx, (255,255,255), cv2.FILLED)
        # cv2.imshow('test', mask)
        object_class_name = 'building'
        output_filename = '{}_{}_{}.jpg'.format(image_id, object_class_name, idx)
        cv2.imwrite(sep_mask_folder + output_filename, mask)


mask_folder = 'D:/Programs/mc_data/test/masks/'
sep_mask_folder = 'D:/Programs/mc_data/test/sep_masks/'
for img in os.listdir(mask_folder):
    filename = mask_folder + img
    image_id = int(img.split('.')[0])
    separate_mask(image_id, filename)

